use std::path::PathBuf;

use futures_util::StreamExt;
use indicatif::{ProgressBar, ProgressStyle};
use tokio::fs::File;
use tokio::io::AsyncWriteExt;

use crate::config::{Model, SDConfig};
use crate::Result;

pub async fn download_model(url: String, output: String, name: String) -> Result<()> {
    let response = reqwest::get(url).await?;
    let content_length = response.content_length().unwrap();
    let mut stream = response.bytes_stream();
    let mut output_file = File::create(&output.clone()).await?;
    let pb = ProgressBar::new(content_length);
    pb.set_style(
        ProgressStyle::default_bar()
            .template("{spinner:.green} [{bar:40.cyan/blue}] {bytes}/{total_bytes} {bytes_per_sec} ({eta}) 【{msg}】")
            .unwrap()
            .progress_chars("#>-"),
    );
    pb.set_message(format!("下载 {name} 中..."));
    while let Some(item) = stream.next().await {
        let chunk = item.unwrap();
        output_file.write_all(&chunk).await?;
        pb.inc(chunk.len() as u64)
    }
    output_file.flush().await?;
    output_file.shutdown().await?;
    pb.finish_with_message(format!("{name} 下载完成!"));

    Ok(())
}

pub async fn download_extra(url: String, output: String) -> Result<()> {
    let response = reqwest::get(url).await?;
    let mut output_file = File::create(&output.clone()).await?;
    let content = response.text().await?;
    output_file.write_all(content.as_bytes()).await?;
    output_file.flush().await?;
    output_file.shutdown().await?;
    Ok(())
}

pub async fn download(cfg: &SDConfig) -> Result<()> {
    for m in &cfg.models {
        let root = &cfg.path.clone();
        let mut output_path = PathBuf::new();
        let model_type = &m.model_type;
        let url = m.url.as_str();
        output_path.push(root);
        output_path.push(cfg.model_paths.get(model_type).unwrap());
        output_path.push(m.get_model_name());
        download_model(
            url.to_string(),
            output_path.to_str().unwrap().to_string(),
            m.get_model_name(),
        )
        .await?;
        if let Some(extra) = &m.extra {
            let mut extra_output_path = PathBuf::new();
            let model_type = &m.model_type;
            extra_output_path.push(root);
            extra_output_path.push(cfg.model_paths.get(model_type).unwrap());
            extra_output_path.push(m.get_extra_name().unwrap());
            download_extra(
                extra.to_owned(),
                extra_output_path.to_str().unwrap().to_string(),
            )
            .await?;
        }
    }
    Ok(())
}
